package com.example.yingjiguanli_demo.interceptor;

import com.example.yingjiguanli_demo.utils.UserUtils;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.Map;
import java.util.Objects;

public class CustomHandshakeInterceptor implements HandshakeInterceptor {

    private final UserUtils userUtils;

    public CustomHandshakeInterceptor(UserUtils userUtils) {
        this.userUtils = userUtils;
    }

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        // 从请求中获取 SESSIONID
        String sessionId = null;
        if (request.getHeaders().get("Cookie") != null) {
            for (String cookie : Objects.requireNonNull(request.getHeaders().get("Cookie"))) {
                if (cookie.startsWith("SESSIONID=")) {
                    sessionId = cookie.substring("SESSIONID=".length());
                    break;
                }
            }
        }

        if (sessionId != null) {
            // 通过 SESSIONID 获取用户 ID
            Long userId = (long) userUtils.getUserId(sessionId);
            // 将 userId 放入 WebSocketSession 的 attributes 中
            attributes.put("userId", userId);
            System.out.println("握手成功，用户ID: " + userId);
            return true; // 返回 true 继续握手
        }

        // 如果没有找到 userId，拒绝连接
        System.out.println("握手失败，未找到有效的 SESSIONID 或用户ID");
        return false;
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                               WebSocketHandler wsHandler, Exception exception) {
        // 可以在握手后做一些额外的处理
    }
}
